import matplotlib.pyplot as plt

epochs = [i + 1 for i in range(10)]
sgd = [0.261355, 0.064336, 0.047649, 0.037445, 0.030153,
       0.026014, 0.023644, 0.019509, 0.016979, 0.01420]
dpsgd = [1.538121, 0.842057, 0.753780, 0.664425, 0.594647,
         0.549227, 0.520469, 0.514079, 0.503336, 0.499845]
loss_no_noise = [2.246184, 1.863532, 1.592694, 1.234314, 0.919571,
                 0.810837, 0.728510, 0.534165, 0.412636, 0.378559]
loss_dp_noise = [2.251099, 2.036245, 1.709324, 1.482202, 1.359249,
                 1.274870, 1.153332, 1.031263, 0.949670, 0.924840]
loss_mdp_noise = [2.281124, 2.159963, 1.972602, 1.831085, 1.735435,
                  1.666723, 1.625087, 1.586538, 1.525584, 1.450791]
loss_cov = [2.294683, 2.253925, 2.153341, 1.984100, 1.826379,
            1.648866, 1.509856, 1.386337, 1.201994, 1.060435]
loss_std = [2.273022, 2.082790, 1.784996, 1.544979, 1.377467,
            1.198416, 1.098116, 1.034869, 1.000428, 0.973738]
loss_both = [2.292810, 2.214785, 2.045617, 1.839140, 1.687974,
             1.591915, 1.505615, 1.396246, 1.265391, 1.109026]

plt.xlabel('epoch')
plt.ylabel('loss')

# plt.plot(epochs, sgd, '-o')
# plt.plot(epochs, dpsgd, '-^')
# plt.legend(['SGD', 'DPSGD'])

# plt.plot(epochs, loss_no_noise, '-o')
# plt.plot(epochs, loss_dp_noise, '-^')
# plt.plot(epochs, loss_mdp_noise, '-+')
# plt.legend(['no noise', 'dp noise', 'mdp noise'])

plt.plot(epochs, loss_mdp_noise, '-+')
plt.plot(epochs, loss_cov, '-s')
plt.plot(epochs, loss_std, '-v')
plt.plot(epochs, loss_both, '-d')
plt.legend(['fixed cov and std', 'adaptive cov', 'adaptive std', 'adaptive cov and std'])

plt.show()
